/*
 * @Author: xiangru.xiao xiangru.xiao@mthreads.com
 * @Date: 2023-05-25 09:32:54
 * @LastEditors: xiangru.xiao xiangru.xiao@mthreads.com
 * @LastEditTime: 2023-05-25 17:41:03
 * Copyright (c) 2023 by Moore Threads, All Rights Reserved. 
 */

#include <string.h>
#include <vector>
#include "objectfactory.h"
#include "node.h"
#include "stopwatch.h"
#include "clcheck.h"
#include "cldeviceinfo.h"
#include "clevent.h"
#include "clFFT.h"
#include "math.h"

#define EPS 1e-6
#define TEST_ITERS 10
struct params
{
    int length;
    int batch;
};

class FFT_test : public Node
{
public:
    FFT_test()
    {
        nodeUint = "GBPS";
    }
    std::string name() override
    {
        return "fft-test";
    }
    std::string describe() override
    {
        return "快速傅里叶变换";
    }
    Node::NodeType  nodeType() override
    {
        return NodeTP_MemRw;
    }
protected:
    bool startup() override;
    bool perform() override;
    void shutdown() override;
private:
    std::vector<params> paramList() const;
    bool testResult(clfftPlanHandle plan, float* h_x, cl_mem d_x, int length, int batch);
private:
    struct cl_handle_t {
        cl_device_id        deviceid;
        cl_context          context;
        cl_command_queue    queue;
    } handle;
};

bool FFT_test::startup()
{
    memset(&handle, 0, sizeof(cl_handle_t));

    handle.deviceid = deviceinfo->deviceID();
    handle.context = clCreateContext(NULL, 1,  &handle.deviceid, NULL, NULL, NULL);
    handle.queue = clCreateCommandQueue(handle.context, handle.deviceid,
                                        CL_QUEUE_PROFILING_ENABLE,
                                        NULL);
    return true;
}

bool FFT_test::perform()
{
    /* Setup clFFT. */
    clfftSetupData fftSetup;
    CHECK_CL_ERROR(clfftInitSetupData(&fftSetup));
    CHECK_CL_ERROR(clfftSetup(&fftSetup));
    float totalValue = 0.0;
    stopwatch timer;
    cl_int err;
    std::vector<params> ls = paramList();
    char* strNum = std::getenv("FFT_TEST_LEVEL");
    int caseNum = ls.size();
    if(strNum)
    {
        caseNum = std::min(caseNum, atoi(strNum));
    }
    for (int i = 0; i < caseNum; ++i) {
        int length = ls[i].length;
        int batch = ls[i].batch;
        const int complex_bytes = sizeof(float) * 2 * length * batch;
        float *h_x = (float*)malloc(complex_bytes);
        cl_mem d_x = clCreateBuffer(handle.context, CL_MEM_READ_WRITE, complex_bytes, NULL, &err);
        CHECK_CL_ERROR(err);
        /* Create a default plan for a complex FFT. */
        clfftPlanHandle plan;
        size_t lengths[3] = {length, 1, 1};
        CHECK_CL_ERROR(clfftCreateDefaultPlan(&plan, handle.context, CLFFT_1D, lengths));
        /* Set plan parameters. */
        CHECK_CL_ERROR(clfftSetPlanPrecision(plan, CLFFT_SINGLE));
        CHECK_CL_ERROR(clfftSetLayout(plan, CLFFT_COMPLEX_INTERLEAVED,
                            CLFFT_COMPLEX_INTERLEAVED));
        CHECK_CL_ERROR(clfftSetResultLocation(plan, CLFFT_INPLACE));
        CHECK_CL_ERROR(clfftSetPlanBatchSize(plan, batch));
        size_t cst[4] = {1, 0, 0, 0};
        cst[1] = lengths[0] * cst[0];
        cst[2] = lengths[1] * cst[1];
        cst[3] = lengths[2] * cst[2];
        CHECK_CL_ERROR(clfftSetPlanInStride(plan, CLFFT_1D, cst));
        CHECK_CL_ERROR(clfftSetPlanOutStride(plan, CLFFT_1D, cst));
        CHECK_CL_ERROR(clfftSetPlanDistance(plan, cst[3], cst[3]));
        /* Bake the plan. */
        CHECK_CL_ERROR(clfftBakePlan(plan, 1, &(handle.queue), NULL, NULL));
        CHECK_CL_ERROR(clFinish(handle.queue));

        if(testResult(plan, h_x, d_x, length, batch))
        { //result correct and then do profile
            timer.start();
            for (int j = 0; j < TEST_ITERS; ++j) {
                  CHECK_CL_ERROR(clfftEnqueueTransform(plan, CLFFT_FORWARD, 1, &(handle.queue),
                                 0, NULL, NULL, &d_x, NULL, NULL));
                  CHECK_CL_ERROR(clFinish(handle.queue));
            }
            int totalTime  = static_cast<int>(timer.usElapsed());
            CHECK_CL_ERROR(clfftDestroyPlan(&plan));
            CHECK_CL_ERROR(clReleaseMemObject(d_x));
            free(h_x);
            totalValue += 2 * float(complex_bytes) * TEST_ITERS / totalTime / 1e3f;
        }
        else
        { //get wrong result
            nodeValue = 0.0;
            CHECK_CL_ERROR(clfftDestroyPlan(&plan));
            CHECK_CL_ERROR(clReleaseMemObject(d_x));
            free(h_x);
            /* Release clFFT library. */
            clfftTeardown( );
            return false;
        }
        
    }
    nodeValue = totalValue;
    /* Release clFFT library. */
    clfftTeardown( );
    return true;
}

void FFT_test::shutdown()
{
    if (handle.queue) {
        CHECK_CL_ERROR(clReleaseCommandQueue(handle.queue));
        handle.queue = nullptr;
    }
    if (handle.context) {
        CHECK_CL_ERROR(clReleaseContext(handle.context));
        handle.context = nullptr;
    }
}

std::vector<params> FFT_test::paramList() const
{
    std::vector<params> list;
    params params_[] = {{128, 1<<18},
                        {256, 1<<17},
                        {512, 1<<16},
                        {1024, 1<<15},
                        {2048, 1<<14},
                        };
    const size_t nSizes =  sizeof(params_) / sizeof (params);
    for (int i = 0; i < nSizes; ++i) {
        list.push_back(params_[i]);
    }
    return list;
}

bool FFT_test::testResult(clfftPlanHandle plan, float* h_x, cl_mem d_x, int length, int batch)
{
    int N = length * std::min(batch, 32);
    float* h_out = (float*)malloc(sizeof(float) * 2 * N);
    for(int i = 0; i < N; i++)
    {
        h_x[2 * i] = rand();
        h_x[2 * i + 1] = 0;
    }
    CHECK_CL_ERROR(clEnqueueWriteBuffer(handle.queue, d_x, CL_TRUE, 0, 
                                sizeof(float) * 2 * N, h_x, 0, NULL, NULL));
    CHECK_CL_ERROR(clfftEnqueueTransform(plan, CLFFT_FORWARD, 1, &(handle.queue),
                                 0, NULL, NULL, &d_x, NULL, NULL));
    CHECK_CL_ERROR(clFinish(handle.queue));
    CHECK_CL_ERROR(clfftEnqueueTransform(plan, CLFFT_BACKWARD, 1, &(handle.queue),
                                 0, NULL, NULL, &d_x, NULL, NULL));
    CHECK_CL_ERROR(clFinish(handle.queue));
    CHECK_CL_ERROR(clEnqueueReadBuffer(handle.queue, d_x, CL_TRUE, 0, 
                               sizeof(float) * 2 * N, h_out, 0, NULL, NULL));

    double maxv = 0;
    double nrmse = 0;  // normalized root mean square error
    for (int i = 0; i < N; i++) {
      double dr = h_x[2 * i] - h_out[2 * i];
      double di = h_x[2 * i + 1] - h_out[2 * i + 1];
      maxv = fabs(h_x[2 * i]) > maxv ? fabs(h_x[2 * i]) : maxv;
      maxv = fabs(h_x[2 * i + 1]) > maxv ? fabs(h_x[2 * i + 1]) : maxv;
      nrmse += ((dr * dr) + (di * di));
    }
    nrmse /= (double)(N);
    nrmse = sqrt(nrmse);
    nrmse /= maxv;
    free(h_out);
    return nrmse < EPS;
}

REGISTER_OBJECT(Node, FFT_test)